#include <stdio.h>
#include <stdlib.h>

#include "esp_system.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "dl_tool.hpp"
#include "mnist_model.hpp"

/**
 * @brief Samples in MNIST dataset are repeated in channel to mimic RGB image. 
 * 
 */
__attribute__((aligned(16))) int16_t example_element[] = {0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 3, 3, 3,
                                                          1, 1, 1, 0, 0, 0, 0, 0,
                                                          0, 7, 7, 7, 0, 0, 0, 37,
                                                          37, 37, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 1,
                                                          1, 1, 2, 2, 2, 0, 0, 0,
                                                          27, 27, 27, 84, 84, 84, 11, 11,
                                                          11, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 119, 119, 119, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 1, 1, 1, 0, 0,
                                                          0, 0, 0, 0, 88, 88, 88, 143,
                                                          143, 143, 110, 110, 110, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 22, 22, 22, 93, 93, 93, 106,
                                                          106, 106, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 4,
                                                          4, 4, 0, 0, 0, 53, 53, 53,
                                                          129, 129, 129, 120, 120, 120, 147, 147,
                                                          147, 175, 175, 175, 157, 157, 157, 166,
                                                          166, 166, 135, 135, 135, 154, 154, 154,
                                                          168, 168, 168, 140, 140, 140, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          2, 2, 2, 0, 0, 0, 11, 11,
                                                          11, 137, 137, 137, 130, 130, 130, 128,
                                                          128, 128, 160, 160, 160, 176, 176, 176,
                                                          159, 159, 159, 167, 167, 167, 178, 178,
                                                          178, 149, 149, 149, 151, 151, 151, 144,
                                                          144, 144, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 1, 1, 1, 0, 0, 0,
                                                          2, 2, 2, 1, 1, 1, 0, 0,
                                                          0, 3, 3, 3, 0, 0, 0, 0,
                                                          0, 0, 115, 115, 115, 114, 114, 114,
                                                          106, 106, 106, 137, 137, 137, 168, 168,
                                                          168, 153, 153, 153, 156, 156, 156, 165,
                                                          165, 165, 167, 167, 167, 143, 143, 143,
                                                          157, 157, 157, 158, 158, 158, 11, 11,
                                                          11, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          1, 1, 1, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 3, 3, 3, 0, 0, 0,
                                                          0, 0, 0, 89, 89, 89, 139, 139,
                                                          139, 90, 90, 90, 94, 94, 94, 153,
                                                          153, 153, 149, 149, 149, 131, 131, 131,
                                                          151, 151, 151, 169, 169, 169, 172, 172,
                                                          172, 143, 143, 143, 159, 159, 159, 169,
                                                          169, 169, 48, 48, 48, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 2, 2, 2, 4, 4, 4,
                                                          1, 1, 1, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 98, 98, 98, 136,
                                                          136, 136, 110, 110, 110, 109, 109, 109,
                                                          110, 110, 110, 162, 162, 162, 135, 135,
                                                          135, 144, 144, 144, 149, 149, 149, 159,
                                                          159, 159, 167, 167, 167, 144, 144, 144,
                                                          158, 158, 158, 169, 169, 169, 119, 119,
                                                          119, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 2, 2, 2, 2, 2, 2,
                                                          1, 1, 1, 2, 2, 2, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 26, 26, 26, 108, 108, 108,
                                                          117, 117, 117, 99, 99, 99, 111, 111,
                                                          111, 117, 117, 117, 136, 136, 136, 156,
                                                          156, 156, 134, 134, 134, 154, 154, 154,
                                                          154, 154, 154, 156, 156, 156, 160, 160,
                                                          160, 141, 141, 141, 147, 147, 147, 156,
                                                          156, 156, 178, 178, 178, 0, 0, 0,
                                                          3, 3, 3, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 21, 21, 21,
                                                          53, 53, 53, 92, 92, 92, 117, 117,
                                                          117, 111, 111, 111, 103, 103, 103, 115,
                                                          115, 115, 129, 129, 129, 134, 134, 134,
                                                          143, 143, 143, 154, 154, 154, 165, 165,
                                                          165, 170, 170, 170, 154, 154, 154, 151,
                                                          151, 151, 154, 154, 154, 143, 143, 143,
                                                          138, 138, 138, 150, 150, 150, 165, 165,
                                                          165, 43, 43, 43, 0, 0, 0, 0,
                                                          0, 0, 23, 23, 23, 54, 54, 54,
                                                          65, 65, 65, 76, 76, 76, 85, 85,
                                                          85, 118, 118, 118, 128, 128, 128, 123,
                                                          123, 123, 111, 111, 111, 113, 113, 113,
                                                          118, 118, 118, 127, 127, 127, 125, 125,
                                                          125, 139, 139, 139, 133, 133, 133, 136,
                                                          136, 136, 160, 160, 160, 140, 140, 140,
                                                          155, 155, 155, 161, 161, 161, 144, 144,
                                                          144, 155, 155, 155, 172, 172, 172, 161,
                                                          161, 161, 189, 189, 189, 62, 62, 62,
                                                          0, 0, 0, 68, 68, 68, 94, 94,
                                                          94, 90, 90, 90, 111, 111, 111, 114,
                                                          114, 114, 111, 111, 111, 114, 114, 114,
                                                          115, 115, 115, 127, 127, 127, 135, 135,
                                                          135, 136, 136, 136, 143, 143, 143, 126,
                                                          126, 126, 127, 127, 127, 151, 151, 151,
                                                          154, 154, 154, 143, 143, 143, 148, 148,
                                                          148, 125, 125, 125, 162, 162, 162, 162,
                                                          162, 162, 144, 144, 144, 138, 138, 138,
                                                          153, 153, 153, 162, 162, 162, 196, 196,
                                                          196, 58, 58, 58, 70, 70, 70, 169,
                                                          169, 169, 129, 129, 129, 104, 104, 104,
                                                          98, 98, 98, 100, 100, 100, 94, 94,
                                                          94, 97, 97, 97, 98, 98, 98, 102,
                                                          102, 102, 108, 108, 108, 106, 106, 106,
                                                          119, 119, 119, 120, 120, 120, 129, 129,
                                                          129, 149, 149, 149, 156, 156, 156, 167,
                                                          167, 167, 190, 190, 190, 190, 190, 190,
                                                          196, 196, 196, 198, 198, 198, 198, 198,
                                                          198, 187, 187, 187, 197, 197, 197, 189,
                                                          189, 189, 184, 184, 184, 36, 36, 36,
                                                          16, 16, 16, 126, 126, 126, 171, 171,
                                                          171, 188, 188, 188, 188, 188, 188, 184,
                                                          184, 184, 171, 171, 171, 153, 153, 153,
                                                          135, 135, 135, 120, 120, 120, 126, 126,
                                                          126, 127, 127, 127, 146, 146, 146, 185,
                                                          185, 185, 195, 195, 195, 209, 209, 209,
                                                          208, 208, 208, 255, 255, 255, 209, 209,
                                                          209, 177, 177, 177, 245, 245, 245, 252,
                                                          252, 252, 251, 251, 251, 251, 251, 251,
                                                          247, 247, 247, 220, 220, 220, 206, 206,
                                                          206, 49, 49, 49, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 12, 12, 12,
                                                          67, 67, 67, 106, 106, 106, 164, 164,
                                                          164, 185, 185, 185, 199, 199, 199, 210,
                                                          210, 210, 211, 211, 211, 210, 210, 210,
                                                          208, 208, 208, 190, 190, 190, 150, 150,
                                                          150, 82, 82, 82, 8, 8, 8, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          178, 178, 178, 208, 208, 208, 188, 188,
                                                          188, 175, 175, 175, 162, 162, 162, 158,
                                                          158, 158, 151, 151, 151, 11, 11, 11,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0,
                                                          0, 0, 0, 0, 0, 0, 0, 0};

extern "C" void app_main(void)
{
    // input
    Tensor<int16_t> input;
    input.set_element((int16_t *)example_element).set_exponent(0).set_shape({28, 28, 3}).set_auto_free(false);

    MNIST model;

    dl::tool::Latency latency;

    // model forward
    latency.start();
    model.forward(input);
    latency.end();
    latency.print("MNIST", "forward");

    // parse
    int16_t *score = model.l5_compress.get_output().get_element_ptr();
    int16_t max_score = score[0];
    int max_index = 0;
    printf("%d, ", max_score);

    for (size_t i = 1; i < 10; i++)
    {
        printf("%d, ", score[i]);
        if (score[i] > max_score)
        {
            max_score = score[i];
            max_index = i;
        }
    }
    printf("\nPrediction Result: %d\n", max_index);
    // PC
    // -7175, -9797, -12315, -11419, -12361, -1369, -11728, -113, -11453, 7859
    // Prediction Result: 9

    // esp32
    // -7170, -9792, -12301, -11416, -12349, -1350, -11715, -118, -11433, 7856
    // Prediction Result: 9

    // esp32s2
    // -7170, -9792, -12301, -11416, -12349, -1350, -11715, -118, -11433, 7856
    // Prediction Result: 9

    // esp32s3
    // -7170, -9792, -12301, -11416, -12349, -1350, -11715, -118, -11433, 7856
    // Prediction Result: 9

    // esp32c3
    // -7170, -9792, -12301, -11416, -12349, -1350, -11715, -118, -11433, 7856
    // Prediction Result: 9
}